#!/usr/bin/env python
"""
IndexTTS2 极简 API - 一个接口搞定所有
"""

import base64
import tempfile
from pathlib import Path
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
from indextts.infer_v2 import IndexTTS2
# 导入 CORSMiddleware
from fastapi.middleware.cors import CORSMiddleware

# 初始化
app = FastAPI(title="IndexTTS2 API")

# --- 新增：配置CORS ---
# 允许所有来源的跨域请求
origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    # 明确允许 POST 和用于预检的 OPTIONS 方法
    allow_methods=["POST", "OPTIONS"],
    allow_headers=["*"],  # 允许所有HTTP头
)
# ----------------------

# 加载模型
tts = IndexTTS2(
    cfg_path="/root/models/IndexTTS-2/config.yaml",
    model_dir="/root/models/IndexTTS-2/",
    use_fp16=False,
    use_cuda_kernel=False,
    use_deepspeed=False
)

class TTSRequest(BaseModel):
    text: str
    voice_base64: str
    # 情绪控制（可选）
    emo_text: str = None  # 情绪文本描述
    emo_audio_base64: str = None  # 情绪参考音频
    emotion_vector: list = None  # 8维情绪向量 [开心,愤怒,悲伤,恐惧,厌恶,忧郁,惊讶,平静]
    emo_alpha: float = 0.7  # 情绪强度
    # 时长控制（可选）
    token_count: int = None  # 精确控制token数
    # 其他
    use_random: bool = False


@app.post("/tts")
async def generate_speech(req: TTSRequest):
    """
    语音合成接口
    """

    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)

        # 保存音色参考
        voice_path = tmpdir / "voice.wav"
        voice_path.write_bytes(base64.b64decode(req.voice_base64))

        # 输出路径
        output_path = tmpdir / "output.wav"

        # 构建参数
        kwargs = {
            "spk_audio_prompt": str(voice_path),
            "text": req.text,
            "output_path": str(output_path),
            "emo_alpha": req.emo_alpha,
            "use_random": req.use_random,
            "verbose": False
        }

        # 情绪控制
        if req.emo_text:
            kwargs["use_emo_text"] = True
            kwargs["emo_text"] = req.emo_text
        elif req.emo_audio_base64:
            emo_path = tmpdir / "emotion.wav"
            emo_path.write_bytes(base64.b64decode(req.emo_audio_base64))
            kwargs["emo_audio_prompt"] = str(emo_path)
        elif req.emotion_vector:
            kwargs["emotion_vector"] = req.emotion_vector

        # 时长控制
        if req.token_count:
            kwargs["token_count"] = req.token_count

        # 执行合成
        try:
            tts.infer(**kwargs)

            # 返回结果
            audio_bytes = output_path.read_bytes()
            return {
                "audio_base64": base64.b64encode(audio_bytes).decode(),
                "success": True
            }
        except Exception as e:
            raise HTTPException(500, f"合成失败: {str(e)}")


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=6006)